// SPDX-License-Identifier: GPL-2.0
/*
 * lombo_mtdnor.c - MTD SPI Driver for Nor
 *
 * Copyright (C) 2016-2021, LomboTech Co.Ltd.
 * Author: lomboswer <lomboswer@lombotech.com>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

#include <linux/init.h>
#include <linux/err.h>
#include <linux/errno.h>
#include <linux/module.h>
#include <linux/device.h>
#include <linux/interrupt.h>
#include <linux/mutex.h>
#include <linux/math64.h>
#include <linux/slab.h>
#include <linux/sched.h>
#include <linux/mod_devicetable.h>

#include <linux/mtd/cfi.h>
#include <linux/mtd/mtd.h>
#include <linux/mtd/partitions.h>
#include <linux/of_platform.h>

#include <linux/spi/spi.h>
#include <linux/spi/flash.h>

#include <mach/debug.h>
#include "lombo_mtdnor.h"

#undef MOD_NAME
#define MOD_NAME "MTDNOR"

#define DEFAULT_SPEED_HZ	(50000000)
#define SIZE_4KBYTE		(4 * 1024)
#define SIZE_16MB		(0x1000000)
#define DEFAULT_RWLINE		(4)

static const char * const probes[] = { "cmdlinepart", NULL };

/**
 * __read_status_cmd - read nor status without address
 * @spi: pointer to spi device.
 * @cmd: control command to send
 * @buf: buffer for read
 * @n_byte: number of bytes to read
 *
 * Return 0 if success, !0 error.
 */
static int __read_status_cmd(struct spi_device *spi,
				u8 cmd, u8 *buf, u32 n_byte)
{
	int err = 0;
	struct spi_transfer t;
	struct spi_message m;
	struct lombo_spi_operate_para esop;

	memset(&t, 0, sizeof(t));
	memset(&esop, 0, sizeof(esop));
	spi_message_init(&m);

	esop.inst = cmd;
	esop.ff = FRAME_FORMAT_STANDARD_SPI;

	t.bits_per_word = 8;
	t.len = n_byte;
	t.rx_buf = buf;
	t.esop = &esop;
	spi_message_add_tail(&t, &m);

	err = spi_sync(spi, &m);
	if (err < 0) {
		PRT_ERR("read status cmd:0x%x failed\n", cmd);
		return err;
	}

	return err;
}

/**
 * __weite_status_cmd - write spi status without address
 * @spi: pointer to spi device.
 * @cmd: control command to send.
 * @buf: buffer for write
 * @n_byte: number of bytes to read
 *
 * Return 0 if success, !0 error.
 */

static int __write_status_cmd(struct spi_device *spi, u8 cmd, u8 *buf, u32 n_byte)
{
	int err = 0;
	struct spi_transfer t;
	struct spi_message m;
	struct lombo_spi_operate_para esop;

	memset(&t, 0, sizeof(t));
	memset(&esop, 0, sizeof(esop));
	spi_message_init(&m);

	esop.inst = cmd;
	esop.ff = FRAME_FORMAT_STANDARD_SPI;

	t.bits_per_word = 8;
	t.len = n_byte;
	t.tx_buf = buf;
	t.esop = &esop;
	spi_message_add_tail(&t, &m);

	err = spi_sync(spi, &m);
	if (err < 0) {
		PRT_ERR("write status cmd: 0x%x failed\n", cmd);
		return err;
	}

	return err;
}

/**
 * __send_ctrl_cmd - send control command
 * @spi: pointer to spi device.
 * @cmd: control command to send
 *
 * Return 0 if success, !0 error.
 */
static int __send_ctrl_cmd(struct spi_device *spi, u8 cmd)
{
	int err = 0;
	struct spi_transfer t;
	struct spi_message m;
	struct lombo_spi_operate_para esop;

	memset(&t, 0, sizeof(t));
	memset(&esop, 0, sizeof(esop));
	spi_message_init(&m);

	esop.inst = cmd;
	esop.ff = FRAME_FORMAT_STANDARD_SPI;

	t.bits_per_word = 8;
	t.esop = &esop;
	spi_message_add_tail(&t, &m);

	err = spi_sync(spi, &m);
	if (err < 0) {
		PRT_ERR("send control cmd:0x%x failed\n", cmd);
		return err;
	}

	return err;
}

/**
 * __send_erase_cmd - send erase command
 * @spi: pointer to spi device.
 * @cmd: erase command to send
 * @addr: nor address for erase
 *
 * Return 0 if success, !0 error.
 */
static int __send_erase_cmd(struct spi_device *spi, u8 cmd, u32 addr)
{
	int err = 0;
	struct spi_transfer t;
	struct spi_message m;
	struct lombo_spi_operate_para esop;
	struct lombo_nor *flash = dev_get_drvdata(&spi->dev);

	memset(&t, 0, sizeof(t));
	memset(&esop, 0, sizeof(esop));
	spi_message_init(&m);

	if (flash->addr_width == WIDTH_3BYTE)
		esop.addr_bit_len = 24;
	else if (flash->addr_width == WIDTH_4BYTE)
		esop.addr_bit_len = 32;

	esop.inst = cmd;
	esop.ff = FRAME_FORMAT_STANDARD_SPI;
	esop.addr[0] = addr;

	t.bits_per_word = 8;
	t.esop = &esop;
	spi_message_add_tail(&t, &m);

	err = spi_sync(spi, &m);
	if (err < 0) {
		PRT_ERR("send erase cmd:0x%x failed\n", cmd);
		return err;
	}

	return err;
}

/**
 * __read_buf_32 - read chip data into buffer in 32 bits read.
 * @spi: pointer to spi device.
 * @addr: nor address for data read
 * @buf: buffer to store data
 * @len: number of bytes to read
 *
 * Return 0 if success, !0 error.
 */
static int __read_buf_32(struct spi_device *spi, u32 addr, size_t buf, u32 len)
{
	int err = 0;
	struct spi_transfer t;
	struct spi_message m;
	struct lombo_spi_operate_para esop;
	struct lombo_nor *flash = dev_get_drvdata(&spi->dev);

	memset(&t, 0, sizeof(t));
	memset(&esop, 0, sizeof(esop));
	spi_message_init(&m);

	if (flash->addr_width == WIDTH_3BYTE) {
		 /* addr 3byte branch */
		if (flash->read_line == 1) {
			esop.inst = OPCODE_NORM_READ;
			esop.ff = FRAME_FORMAT_STANDARD_SPI;
		} else if (flash->read_line == 2) {
			esop.inst = OPCODE_DUAL_READ;
			esop.ff = FRAME_FORMAT_DUAL_SPI;
			esop.wait_cycle = 8;
		} else if (flash->read_line == 4) {
			esop.inst = OPCODE_QUAD_READ;
			esop.ff = FRAME_FORMAT_QUAD_SPI;
			esop.wait_cycle = 8;
		}

		esop.addr_bit_len = 24;
	} else if (flash->addr_width == WIDTH_4BYTE) {
		if (flash->read_line == 1) {
			esop.inst = OPCODE_NORM_READ4B;
			esop.ff = FRAME_FORMAT_STANDARD_SPI;
		} else if (flash->read_line == 2) {
			esop.inst = OPCODE_DUAL_READ4B;
			esop.ff = FRAME_FORMAT_DUAL_SPI;
			esop.wait_cycle = 8;
		} else if (flash->read_line == 4) {
			esop.inst = OPCODE_QUAD_READ4B;
			esop.ff = FRAME_FORMAT_QUAD_SPI;
			esop.wait_cycle = 8;
		}

		esop.addr_bit_len = 32;
	}

	esop.addr[0] = addr;

	t.speed_hz = flash->freq;
	t.bits_per_word = 32;
	t.len = len;
	t.rx_buf = (u32 *)buf;
	t.esop = &esop;
	spi_message_add_tail(&t, &m);

	err = spi_sync(spi, &m);
	if (err < 0) {
		PRT_ERR("read from:0x%x to:0x%px, len:%d failed\n",
				addr, (size_t *)buf, len);
		return err;
	}

	return err;
}

/**
 * __read_buf_8 - read chip data into buffer in 8 bits read.
 * @spi: pointer to spi device.
 * @addr: nor address for data read
 * @buf: buffer to store data
 * @len: number of bytes to read
 *
 * Return 0 if success, !0 error.
 */
static int __read_buf_8(struct spi_device *spi, u32 addr, size_t buf, u32 len)
{
	int err = 0;
	struct spi_transfer t;
	struct spi_message m;
	struct lombo_spi_operate_para esop;
	struct lombo_nor *flash = dev_get_drvdata(&spi->dev);

	memset(&t, 0, sizeof(t));
	memset(&esop, 0, sizeof(esop));
	spi_message_init(&m);

	if (flash->addr_width == WIDTH_3BYTE) {
		/* addr 3byte branch */
		if (flash->read_line == 1) {
			esop.inst = OPCODE_NORM_READ;
			esop.ff = FRAME_FORMAT_STANDARD_SPI;
		} else if (flash->read_line == 2) {
			esop.inst = OPCODE_DUAL_READ;
			esop.ff = FRAME_FORMAT_DUAL_SPI;
			esop.wait_cycle = 8;
		} else if (flash->read_line == 4) {
			esop.inst = OPCODE_QUAD_READ;
			esop.ff = FRAME_FORMAT_QUAD_SPI;
			esop.wait_cycle = 8;
		}

		esop.addr_bit_len = 24;
	} else if (flash->addr_width == WIDTH_4BYTE) {
		/* addr 4byte branch */
		if (flash->read_line == 1) {
			esop.inst = OPCODE_NORM_READ4B;
			esop.ff = FRAME_FORMAT_STANDARD_SPI;
		} else if (flash->read_line == 2) {
			esop.inst = OPCODE_DUAL_READ4B;
			esop.ff = FRAME_FORMAT_DUAL_SPI;
			esop.wait_cycle = 8;
		} else if (flash->read_line == 4) {
			esop.inst = OPCODE_QUAD_READ4B;
			esop.ff = FRAME_FORMAT_QUAD_SPI;
			esop.wait_cycle = 8;
		}

		esop.addr_bit_len = 32;
	}

	esop.addr[0] = addr;

	t.speed_hz = flash->freq;
	t.bits_per_word = 8;
	t.len = len;
	t.rx_buf = (u8 *)buf;
	t.esop = &esop;
	spi_message_add_tail(&t, &m);

	err = spi_sync(spi, &m);
	if (err < 0) {
		PRT_ERR("read from:0x%x to:0x%px, len:%d failed\n",
				addr, (size_t *)buf, len);
		return err;
	}

	return err;
}

/**
 * __read_buf - read chip data into buffer
 * @spi: pointer to spi device.
 * @addr: nor address for data read
 * @buf: buffer to store data
 * @len: number of bytes to read
 *
 * Return 0 if success, !0 error.
 */
static int __read_buf(struct spi_device *spi, u32 addr, size_t buf, u32 len)
{
	int err = 0;
	u32 nor_addr = addr;
	size_t dst_buf = buf;
	int div_len = len / 4;
	int mod_len = len & (4 - 1);

	if (addr > 0 && (addr % 4) != 0)
		PRT_ERR("addr=0x%x, len=%d\n", addr, len);

	if (div_len > 0) {
		/* read nor to buffer */
		err = __read_buf_32(spi, nor_addr, dst_buf, 4 * div_len);
		if (err) {
			PRT_ERR("__read_buf_32 err, ret=%d\n", err);
			err = __LINE__;
			goto out;
		}

		nor_addr += 4 * div_len;
		dst_buf  += 4 * div_len;
	}

	if (mod_len != 0) {
		/* read nor to buffer */
		err = __read_buf_8(spi, nor_addr, dst_buf, mod_len);
		if (err) {
			PRT_ERR("__read_buf_8 err, ret=%d\n", err);
			err = __LINE__;
			goto out;
		}
	}

out:
	return err;
}

/**
 * lombo_nor_read_buf - read chip data into buffer
 * @spi: pointer to spi device.
 * @addr: nor address for data read
 * @buf: buffer to store data
 * @len: number of bytes to read
 *
 * Return 0 if success, !0 error.
 */
static int lombo_nor_read_buf(struct spi_device *spi,
				u32 addr, u8 *buf, u32 len)
{
	int err = 0;
	u32 nor_addr = addr;
	size_t dst_buf = (size_t)buf;
	int div_len = len / (BUF_READ_SIZE);
	int mod_len = len & (BUF_READ_SIZE - 1);
	int i;

	for (i = 0; i < div_len; i++) {
		/* read nor to buffer */
		err = __read_buf(spi, nor_addr, dst_buf, BUF_READ_SIZE);
		if (err) {
			PRT_ERR("__read_buf err, ret=%d\n", err);
			err = __LINE__;
			goto out;
		}

		nor_addr += BUF_READ_SIZE;
		dst_buf  += BUF_READ_SIZE;
	}

	if (mod_len != 0) {
		/* read nor to buffer */
		err = __read_buf(spi, nor_addr, dst_buf, mod_len);
		if (err) {
			PRT_ERR("__read_buf err, ret=%d\n", err);
			err = __LINE__;
			goto out;
		}
	}

out:
	return err;
}

/**
 * __write_buf_32 - write buffer to chip in 32 bits write.
 * @spi: pointer to spi device.
 * @addr: nor address for data write
 * @buf: data buffer
 * @len: number of bytes to write
 *
 * Return 0 if success, !0 error.
 */
static int __write_buf_32(struct spi_device *spi,
				u32 addr, size_t buf, u32 len)
{
	int err = 0;
	struct spi_transfer t;
	struct spi_message m;
	struct lombo_spi_operate_para esop;
	struct lombo_nor *flash = dev_get_drvdata(&spi->dev);

	memset(&t, 0, sizeof(t));
	memset(&esop, 0, sizeof(esop));
	spi_message_init(&m);

	if (flash->addr_width == WIDTH_3BYTE) {
		/* 3byte addr branch */
		if (flash->write_line == 1) {
			esop.inst = OPCODE_PP;
			esop.ff = FRAME_FORMAT_STANDARD_SPI;
		} else if (flash->write_line == 4) {
			esop.inst = OPCODE_QUAD_PP;
			esop.ff = FRAME_FORMAT_QUAD_SPI;
		}

		esop.addr_bit_len = 24;
	} else if (flash->addr_width == WIDTH_4BYTE) {
		/* 4byte addr branch */
		if (flash->write_line == 1) {
			esop.inst = OPCODE_PP4B;
			esop.ff = FRAME_FORMAT_STANDARD_SPI;
		} else if (flash->write_line == 4) {
			esop.inst = OPCODE_QUAD_PP4B;
			esop.ff = FRAME_FORMAT_QUAD_SPI;
		}

		esop.addr_bit_len = 32;
	}

	esop.addr[0] = addr;

	t.speed_hz = flash->freq;
	t.bits_per_word = 32;
	t.len = len;
	t.tx_buf = (u32 *)buf;
	t.esop = &esop;
	spi_message_add_tail(&t, &m);

	err = spi_sync(spi, &m);
	if (err < 0) {
		PRT_ERR("write from:0x%px to:0x%x, len:%d failed\n",
				(size_t *)buf, addr, len);
		return err;
	}

	return err;
}

/**
 * __write_buf_8 - write buffer to chip in 8 bits write.
 * @spi: pointer to spi device.
 * @addr: nor address for data write
 * @buf: data buffer
 * @len: number of bytes to write
 *
 * Return 0 if success, !0 error.
 */
static int __write_buf_8(struct spi_device *spi, u32 addr, size_t buf, u32 len)
{
	int err = 0;
	struct spi_transfer t;
	struct spi_message m;
	struct lombo_spi_operate_para esop;
	struct lombo_nor *flash = dev_get_drvdata(&spi->dev);

	memset(&t, 0, sizeof(t));
	memset(&esop, 0, sizeof(esop));
	spi_message_init(&m);

	if (flash->addr_width == WIDTH_3BYTE) {
		/* 3byte addr branch */
		if (flash->write_line == 1) {
			esop.inst = OPCODE_PP;
			esop.ff = FRAME_FORMAT_STANDARD_SPI;
		} else if (flash->write_line == 4) {
			esop.inst = OPCODE_QUAD_PP;
			esop.ff = FRAME_FORMAT_QUAD_SPI;
		}

		esop.addr_bit_len = 24;
	} else if (flash->addr_width == WIDTH_4BYTE) {
		/* 4byte branch */
		if (flash->write_line == 1) {
			esop.inst = OPCODE_PP4B;
			esop.ff = FRAME_FORMAT_STANDARD_SPI;
		} else if (flash->write_line == 4) {
			esop.inst = OPCODE_QUAD_PP4B;
			esop.ff = FRAME_FORMAT_QUAD_SPI;
		}

		esop.addr_bit_len = 32;
	}

	esop.addr[0] = addr;

	t.speed_hz = flash->freq;
	t.bits_per_word = 8;
	t.len = len;
	t.tx_buf = (u8 *)buf;
	t.esop = &esop;

	spi_message_add_tail(&t, &m);

	err = spi_sync(spi, &m);
	if (err < 0) {
		PRT_ERR("write from:0x%px to:0x%x, len:%d failed\n",
				(size_t *)buf, addr, len);
		return err;
	}

	return err;
}

/**
 * lombo_read_sr - read the status register.
 * @flash: pointer to lombo nor.
 *
 * Return the status register value if success, <0 error.
 */
static int lombo_read_sr(struct lombo_nor *flash, u8 cmd)
{
	int err = 0;
	u8 val = 0;

	err = __read_status_cmd(flash->spi, cmd, &val, 1);
	if (err) {
		PRT_ERR("read status register error\n");
		err = -1;
		return err;
	}

	return val;
}

/**
 * lombo_wait_till_ready - read status register until ready, or timeout occurs.
 * @flash: pointer to lombo nor.
 * @status_bit: status bit for check
 * @ready_status: target status, 0 or 1
 *
 * Return 0 if success, !0 error.
 */
static int lombo_wait_till_ready(struct lombo_nor *flash, u8 cmd,
				u8 status_bit, u8 ready_status)
{
	unsigned long deadline;
	int sr;

	deadline = jiffies + MAX_READY_WAIT_JIFFIES;

	do {
		sr = lombo_read_sr(flash, cmd);
		if (sr < 0)
			break;

		if (((sr >> status_bit) & 0x1) == ready_status)
			return 0;

		cond_resched();

	} while (!time_after_eq(jiffies, deadline));

	return 1;
}

/**
 * lombo_write_enable - set write enable latch with Write Enable command.
 * @flash: pointer to lombo nor.
 *
 * Return 0 if success, !0 error.
 */
static int lombo_write_enable(struct lombo_nor *flash)
{
	int err = 0;

	err = __send_ctrl_cmd(flash->spi, OPCODE_WREN);
	if (err) {
		PRT_ERR("write enable error\n");
		goto out;
	}

out:
	return err;
}

static int qe_endisable(struct lombo_nor *flash, u8 cmd, u8 bit_index, int enable)
{
	u8 val;
	int err = 0;

	/* Wait until finished previous write command. */
	if (lombo_wait_till_ready(flash, OPCODE_RDSR1, WIP_BIT, 0)) {
		PRT_ERR("wait WIP timeout\n");
		err = -1;
		return err;
	}

	val = lombo_read_sr(flash, cmd);
	if (val < 0) {
		PRT_ERR("read status register failed\n");
		return val;
	} else if (val & (1 << bit_index))
		return err;

	/* Send write enable, then erase commands. */
	if (lombo_write_enable(flash)) {
		PRT_ERR("write enable failed before erase chip\n");
		err = -1;
		return err;
	}

	if (lombo_wait_till_ready(flash, OPCODE_RDSR1, WEL_BIT, 1)) {
		PRT_ERR("wait WEL timeout\n");
		err = -1;
		return err;
	}

	if (enable)
		val |= (1 << bit_index);
	else
		val &= ~(1 << bit_index);

	if (cmd == OPCODE_RDSR1)
		err = __write_status_cmd(flash->spi, OPCODE_WRSR1, &val, 1);
	else
		err = __write_status_cmd(flash->spi, OPCODE_WRSR2, &val, 1);
	if (err) {
		PRT_ERR("read status register error\n");
		err = -1;
		return err;
	}

	return err;
}

/**
 * lombo_qe_enable - set write enable latch with Write Enable command.
 * @nor: pointer to lombo nor.
 *
 * Return 0 if success, !0 error.
 */
static int lombo_qe_enable(struct lombo_nor *flash)
{
	int err = 0;
	unsigned int id = flash->id;

	if (id == 0xc22018 || id == 0xc22019) {
		err = qe_endisable(flash, OPCODE_RDSR1, QE_BIT6, 1);
		if (err)
			return err;
	} else {
		err = qe_endisable(flash, OPCODE_RDSR2, QE_BIT1, 1);
		if (err)
			return err;
	}

	PRT_DBG("QE BIT has been enabled\n");

	return err;
}

static int lombo_qe_disable(struct lombo_nor *flash)
{
	int err = 0;
	unsigned int id = flash->id;

	if (id == 0xc22018 || id == 0xc22019) {
		err = qe_endisable(flash, OPCODE_RDSR1, QE_BIT6, 0);
		if (err)
			return err;
	} else {
		err = qe_endisable(flash, OPCODE_RDSR2, QE_BIT1, 0);
		if (err)
			return err;
	}

	PRT_DBG("QE BIT has been disabled\n");

	return err;
}

#if 0
int lombo_nor_quad_enable(void)
{
	int err = 0;

	err = lombo_qe_enable(flash);
	if (err) {
		PRT_ERR("enable quda failed\n");
		return err;
	}

	return err;
}

int lombo_nor_quad_disable(void)
{
	int err = 0;

	err = lombo_qe_disable(flash);
	if (err) {
		PRT_ERR("disable quda failed\n");
		return err;
	}

	return err;
}

void lombo_nor_change_readline(int num)
{
	switch (num) {
	case 1:
		flash->read_line = 1;
		break;
	case 2:
		flash->read_line = 2;
		break;
	case 4:
		flash->read_line = 4;
		break;
	default:
		PRT_ERR("change failed!The num(%d) is invalid\n", num);
	}
}

void lombo_nor_change_writeline(int num)
{
	switch (num) {
	case 1:
		flash->write_line = 1;
		break;
	case 2:
		flash->write_line = 2;
		break;
	case 4:
		flash->write_line = 4;
		break;
	default:
		PRT_ERR("Change failed!The num(%d) is invalid\n", num);
	}
}
#endif

static int lombo_nor_4byte_addr_enable(struct lombo_nor *flash)
{
	int err = 0;
	u8 val = 0;

	val = lombo_read_sr(flash, OPCODE_RDCR);
	if (val < 0) {
		PRT_ERR("read configuration register failed\n");
		return val;
	}

	/* when the 4byte addr mode is enable(1), return directly */
	if (val & (1 << CFG_ADDR_4BYTE_BIT)) {
		PRT_DBG("address 4 byte mode already enabled\n");
		return err;
	} else {
		err = __send_ctrl_cmd(flash->spi, OPCODE_EN4B);
		if (err) {
			PRT_ERR("enable 4 byte address functione error\n");
			return err;
		}

		if (lombo_wait_till_ready(flash, OPCODE_RDCR,
					  CFG_ADDR_4BYTE_BIT, 1)) {
			PRT_ERR("wait enable 4 byte address timeout\n");
			err = -1;
			return err;
		}

		PRT_DBG("address 4 byte mode has been enabled\n");
	}

	return err;
}

static int lombo_nor_4byte_addr_disable(struct lombo_nor *flash)
{
	int err = 0;
	u8 val = 0;

	val = lombo_read_sr(flash, OPCODE_RDCR);
	if (val < 0) {
		PRT_ERR("read configuration register failed\n");
		return val;
	}

	if (val & (1 << CFG_ADDR_4BYTE_BIT)) {
		err = __send_ctrl_cmd(flash->spi, OPCODE_EX4B);
		if (err) {
			PRT_ERR("enable 4 byte address functione error\n");
			return err;
		}

		if (lombo_wait_till_ready(flash, OPCODE_RDCR,
					  CFG_ADDR_4BYTE_BIT, 0)) {
			PRT_ERR("wait enable 4 byte address timeout\n");
			err = -1;
			return err;
		}

		PRT_DBG("address 4 byte mode has been disabled\n");
	}

	return err;
}

/**
 * lombo_nor_write_buf - write buffer to chip
 * @spi: pointer to spi device.
 * @addr: nor address for data write
 * @buf: data buffer
 * @len: number of bytes to write
 *
 * Return 0 if success, !0 error.
 */
static int lombo_nor_write_buf(struct spi_device *spi,
				u32 addr, size_t buf, u32 len)
{
	struct lombo_nor *flash = dev_get_drvdata(&spi->dev);
	int err = 0;
	u32 nor_addr = addr;
	size_t src_buf = buf;
	int div_len = len / 4;
	int mod_len = len & (4 - 1);

	if (addr > 0 && (addr % 4) != 0)
		PRT_ERR("addr=0x%x, len=%d\n", addr, len);

	if (div_len > 0) {
		/* write buffer to nor */
		err = __write_buf_32(spi, nor_addr, src_buf, 4 * div_len);
		if (err) {
			PRT_ERR("__write_buf_32 err, ret=%d\n", err);
			err = __LINE__;
			goto out;
		}

		nor_addr += 4 * div_len;
		src_buf  += 4 * div_len;
	}

	if (mod_len != 0) {
		if (div_len > 0) {
			if (lombo_wait_till_ready(flash, OPCODE_RDSR1,
						  WIP_BIT, 0)) {
				PRT_ERR("\n");
				err = __LINE__;
				goto out;
			}

			if (lombo_write_enable(flash)) {
				PRT_ERR("\n");
				err = __LINE__;
				goto out;
			}

			if (lombo_wait_till_ready(flash, OPCODE_RDSR1,
						  WEL_BIT, 1)) {
				PRT_ERR("\n");
				err = __LINE__;
				goto out;
			}
		}

		/* write buffer to nor */
		err = __write_buf_8(spi, nor_addr, src_buf, mod_len);
		if (err) {
			PRT_ERR("__write_buf_8 err, ret=%d\n", err);
			err = __LINE__;
			goto out;
		}
	}

out:
	return err;
}

/**
 * lombo_erase_chip - erase the whole flash memory.
 * @flash: pointer to lombo nor.
 *
 * Return 0 if success, !0 error.
 */
static int lombo_erase_chip(struct lombo_nor *flash)
{
	PRT_DBG("%lldKiB\n", (long long)(flash->mtd.size >> 10));

	/* Wait until finished previous write command. */
	if (lombo_wait_till_ready(flash, OPCODE_RDSR1, WIP_BIT, 0)) {
		PRT_ERR("wait WIP timeout\n");
		return 1;
	}

	/* Send write enable, then erase commands. */
	if (lombo_write_enable(flash)) {
		PRT_ERR("write enable failed before erase chip\n");
		return 1;
	}

	if (lombo_wait_till_ready(flash, OPCODE_RDSR1, WEL_BIT, 1)) {
		PRT_ERR("wait WEL timeout\n");
		return 1;
	}

	return __send_ctrl_cmd(flash->spi, OPCODE_CHIP_ERASE);
}

/**
 * lombo_erase_sector - erase one sector of flash memory at offset.
 * @flash: pointer to lombo nor.
 * @offset: offset in flash.
 *
 * Return 0 if success, !0 error.
 */
static int lombo_erase_sector(struct lombo_nor *flash, u32 offset)
{
	PRT_DBG("%dKiB at 0x%08x\n", flash->mtd.erasesize / 1024, offset);

	/* Wait until finished previous write command. */
	if (lombo_wait_till_ready(flash, OPCODE_RDSR1, WIP_BIT, 0)) {
		PRT_ERR("wait WIP timeout\n");
		return 1;
	}

	/* Send write enable, then erase commands. */
	if (lombo_write_enable(flash)) {
		PRT_ERR("write enable failed before erase sector\n");
		return 1;
	}

	if (lombo_wait_till_ready(flash, OPCODE_RDSR1, WEL_BIT, 1)) {
		PRT_ERR("wait WEL timeout\n");
		return 1;
	}

	return __send_erase_cmd(flash->spi, flash->erase_opcode, offset);
}

/* MTD implementation */

/**
 * lombo_nor_erase - erase an address range on the flash chip.
 * @mtd: pointer to mtd info.
 * @instr: pointer to erase info.
 *
 * Return 0 if success, !0 error.
 */
static int lombo_nor_erase(struct mtd_info *mtd, struct erase_info *instr)
{
	struct lombo_nor *flash = container_of(mtd, struct lombo_nor, mtd);
	u32 addr, len;
	uint32_t rem;

	PRT_DBG("at 0x%llx, len %lld\n", (long long)instr->addr,
				(long long)instr->len);

	div_u64_rem(instr->len, mtd->erasesize, &rem);
	if (rem) {
		PRT_ERR("\n");
		return -EINVAL;
	}

	addr = instr->addr;
	len = instr->len;

	mutex_lock(&flash->lock);

	/* whole-chip erase? */
	if (len == flash->mtd.size) {
		if (lombo_erase_chip(flash)) {
			PRT_ERR("\n");
			mutex_unlock(&flash->lock);
			return -EIO;
		}

	/* REVISIT in some cases we could speed up erasing large regions
	 * by using OPCODE_SE instead of OPCODE_BE_4K.  We may have set up
	 * to use "small sector erase", but that's not always optimal.
	 */

	/* "sector"-at-a-time erase */
	} else {
		while (len) {
			if (lombo_erase_sector(flash, addr)) {
				PRT_ERR("\n");
				mutex_unlock(&flash->lock);
				return -EIO;
			}

			addr += mtd->erasesize;
			len -= mtd->erasesize;
		}
	}

	mutex_unlock(&flash->lock);

	return 0;
}

/**
 * lombo_nor_read - read an address range from the flash chip.
 * @mtd: pointer to mtd info.
 * @from: flash address from read.
 * @len: length for read.
 * @retlen: actual length for read.
 * @buf: buffer for read.
 *
 * Return 0 if success, !0 error.
 */
static int lombo_nor_read(struct mtd_info *mtd, loff_t from, size_t len,
	size_t *retlen, u_char *buf)
{
	struct lombo_nor *flash = container_of(mtd, struct lombo_nor, mtd);
	int err = 0;

	PRT_DBG("from 0x%08x, len %zd\n", (u32)from, len);

	mutex_lock(&flash->lock);

	/* Wait till previous write/erase is done. */
	if (lombo_wait_till_ready(flash, OPCODE_RDSR1, WIP_BIT, 0)) {
		PRT_ERR("\n");
		*retlen = 0;
		err = __LINE__;
		goto out;
	}

	err = lombo_nor_read_buf(flash->spi, from, buf, len);
	if (err) {
		PRT_ERR("\n");
		*retlen = 0;
		err = __LINE__;
		goto out;
	}

	*retlen = len;

out:
	mutex_unlock(&flash->lock);

	return err;
}

/**
 * lombo_nor_write - write an address range to the flash chip.
 * @mtd: pointer to mtd info.
 * @to: flash address to write.
 * @len: length for read.
 * @retlen: actual length for write.
 * @buf: buffer for write.
 *
 * Return 0 if success, !0 error.
 */
static int lombo_nor_write(struct mtd_info *mtd, loff_t to, size_t len,
	size_t *retlen, const u_char *buf)
{
	struct lombo_nor *flash = container_of(mtd, struct lombo_nor, mtd);
	u32 page_offset, page_size;
	int err = 0;

	PRT_DBG("to 0x%08x, len %zd, buf=%p\n", (u32)to, len, buf);

	mutex_lock(&flash->lock);

	/* Wait until finished previous write command. */
	if (lombo_wait_till_ready(flash, OPCODE_RDSR1, WIP_BIT, 0)) {
		PRT_ERR("\n");
		*retlen = 0;
		err = __LINE__;
		goto out;
	}

	if (lombo_write_enable(flash)) {
		PRT_ERR("\n");
		*retlen = 0;
		err = __LINE__;
		goto out;
	}

	if (lombo_wait_till_ready(flash, OPCODE_RDSR1, WEL_BIT, 1)) {
		PRT_ERR("\n");
		*retlen = 0;
		err = __LINE__;
		goto out;
	}

	page_offset = to & (flash->page_size - 1);
	/* do all the bytes fit onto one page? */
	if (page_offset + len <= flash->page_size) {
		err = lombo_nor_write_buf(flash->spi, to, (size_t)buf, len);
		if (err) {
			PRT_ERR("\n");
			*retlen = 0;
			err = __LINE__;
			goto out;
		}

		*retlen = len;
	} else {
		u32 i;

		/* the size of data remaining on the first page */
		page_size = flash->page_size - page_offset;

		err = lombo_nor_write_buf(flash->spi, to,
				(size_t)buf, page_size);
		if (err) {
			PRT_ERR("\n");
			*retlen = 0;
			err = __LINE__;
			goto out;
		}

		*retlen = page_size;

		/* write everything in flash->page_size chunks */
		for (i = page_size; i < len; i += page_size) {
			page_size = len - i;
			if (page_size > flash->page_size)
				page_size = flash->page_size;

			if (lombo_wait_till_ready(flash, OPCODE_RDSR1,
						  WIP_BIT, 0)) {
				PRT_ERR("\n");
				*retlen = 0;
				err = __LINE__;
				goto out;
			}

			if (lombo_write_enable(flash)) {
				PRT_ERR("\n");
				*retlen = 0;
				err = __LINE__;
				goto out;
			}

			if (lombo_wait_till_ready(flash, OPCODE_RDSR1,
						  WEL_BIT, 1)) {
				PRT_ERR("\n");
				*retlen = 0;
				err = __LINE__;
				goto out;
			}

			err = lombo_nor_write_buf(flash->spi, to + i,
						(size_t)buf + i, page_size);
			if (err) {
				PRT_ERR("\n");
				*retlen = 0;
				err = __LINE__;
				goto out;
			}

			*retlen += page_size;
		}
	}

out:
	mutex_unlock(&flash->lock);

	return err;
}

/**
 * lombo_jedec_probe - probe flash by read jedec id.
 * @spi: Pointer to the SPI device.
 *
 * Return the pointer to spi device id if success, otherwise ERR_PTR.
 */
static const struct spi_device_id *lombo_jedec_probe(struct spi_device *spi)
{
	int			tmp;
	u8			code = OPCODE_RDID;
	u8			id[5];
	u32			jedec;
	u16                     ext_jedec;
	struct flash_info	*info;

	/* JEDEC also defines an optional "extended device information"
	 * string for after vendor-specific data, after the three bytes
	 * we use here.  Supporting some chips might require using it.
	 */
	tmp = __read_status_cmd(spi, code, id, 5);
	if (tmp < 0) {
		pr_debug("%s: error %d reading JEDEC ID\n",
				dev_name(&spi->dev), tmp);
		return ERR_PTR(tmp);
	}

	jedec = id[0];
	jedec = jedec << 8;
	jedec |= id[1];
	jedec = jedec << 8;
	jedec |= id[2];

	ext_jedec = id[3] << 8 | id[4];

	for (tmp = 0; tmp < ARRAY_SIZE(lombo_nor_ids) - 1; tmp++) {
		info = (void *)lombo_nor_ids[tmp].driver_data;
		if (info->jedec_id == jedec) {
			if (info->ext_id != 0 && info->ext_id != ext_jedec)
				continue;
			return &lombo_nor_ids[tmp];
		}
	}

	PRT_ERR("unrecognized JEDEC id 0x%06x\n", jedec);

	return ERR_PTR(-ENODEV);
}

/**
 * lombo_nor_probe - probe function of spi nor.
 * @spi: Pointer to the SPI device.
 *
 * Return 0 if success, 0! error.
 */
static int lombo_nor_probe(struct spi_device *spi)
{
	const struct spi_device_id *id = lombo_nor_ids;
	struct lombo_nor *flash;
	struct flash_info *info;
	struct mtd_part_parser_data ppdata;
	struct device_node __maybe_unused *np = spi->dev.of_node;
	const struct spi_device_id *jid;
	int res;

#ifdef CONFIG_MTD_OF_PARTS
	if (!of_device_is_available(np))
		return -ENODEV;
#endif

	info = (void *)id->driver_data;
	if (info->jedec_id) {
		jid = lombo_jedec_probe(spi);
		if (IS_ERR(jid)) {
			PRT_ERR("jedec probe failed\n");
			return PTR_ERR(jid);
		} else if (jid != id) {
			PRT_INFO("found %s\n", jid->name);
			id = jid;
			info = (void *)jid->driver_data;
		}
	}

	flash = kzalloc(sizeof(*flash), GFP_KERNEL);
	if (!flash) {
		PRT_ERR("malloc flash failed\n");
		return -ENOMEM;
	}

	flash->spi = spi;
	flash->size = info->sector_size * info->n_sectors;
	flash->page_size = info->page_size;
	flash->id = info->jedec_id;
	flash->addr_width = WIDTH_3BYTE;

	mutex_init(&flash->lock);
	dev_set_drvdata(&spi->dev, flash);

	flash->mtd.priv = flash;
	flash->mtd.name = "lombo-nor";
	flash->mtd.type = MTD_NORFLASH;
	flash->mtd.writesize = 1;
	flash->mtd.flags = MTD_CAP_NORFLASH;
	flash->mtd.size = info->sector_size * info->n_sectors;
	flash->mtd._erase = lombo_nor_erase;
	flash->mtd._read = lombo_nor_read;
	flash->mtd._write = lombo_nor_write;

	/* prefer "small sector" erase if possible */
	/* if (info->flags & SECT_4K) { */
	if (0) { /* mtdnor always 64k erase */
		if (flash->size > SIZE_16MB)
			flash->erase_opcode = OPCODE_BE4B_4K;
		else
			flash->erase_opcode = OPCODE_BE_4K;

		flash->erase_size = SIZE_4KBYTE;
		flash->mtd.erasesize = SIZE_4KBYTE;
	} else {
		if (flash->size > SIZE_16MB)
			flash->erase_opcode = OPCODE_SE4B;
		else
			flash->erase_opcode = OPCODE_SE;

		flash->erase_size = info->sector_size; /* general 64K bytes */
		flash->mtd.erasesize = info->sector_size;
	}

	res = of_property_read_u32(np, "read-line", &flash->read_line);
	if (res)
		flash->read_line = DEFAULT_RWLINE;

	res = of_property_read_u32(np, "write-line", &flash->write_line);
	if (res)
		flash->write_line = DEFAULT_RWLINE;

	if (flash->read_line == 4 || flash->read_line == 4)
		lombo_qe_enable(flash);

	res = of_property_read_u32(np, "spi-max-frequency", &flash->freq);
	if (res)
		flash->freq = DEFAULT_SPEED_HZ;

	if (flash->size > SIZE_16MB) {
		res = lombo_nor_4byte_addr_enable(flash);
		if (res)
			return -EINVAL;
		flash->addr_width = WIDTH_4BYTE;
	}

	flash->mtd.dev.parent = &spi->dev;
	flash->page_size = info->page_size;
	flash->mtd.writebufsize = flash->page_size;

	mtd_set_of_node(&flash->mtd, np);

	PRT_INFO("%s (%lld Kbytes)\n", id->name,
				(long long)flash->mtd.size >> 10);

	PRT_INFO("mtd name=%s, size=0x%llx(%lldMiB) erasesize=0x%.8x(%uKiB)\n",
		flash->mtd.name,
		(long long)flash->mtd.size, (long long)(flash->mtd.size >> 20),
		flash->mtd.erasesize, flash->mtd.erasesize / 1024);

	if (mtd_device_parse_register(&flash->mtd, probes, &ppdata, NULL, 0)) {
		PRT_ERR("register mtd device failed\n");
		kfree(flash);
		return -EINVAL;
	}

	return 0;
}

/**
 * lombo_nor_remove - remove function of spi nor.
 * @spi: Pointer to the SPI device.
 *
 * Return 0 if success, 0! error.
 */
static int lombo_nor_remove(struct spi_device *spi)
{
	struct lombo_nor *flash = dev_get_drvdata(&spi->dev);
	int status = 0;

	/* Clean up MTD stuff. */
	status = mtd_device_unregister(&flash->mtd);
	if (status == 0)
		kfree(flash);

	return 0;
}

static const struct of_device_id lombo_mtdnor_match[] = {
	{ .compatible = "lombo,nor", },
	{},
};
MODULE_DEVICE_TABLE(of, lombo_nor_match);

static struct spi_driver lombo_mtdnor_driver = {
	.driver = {
		.name  = "lombo_nor",
		.owner = THIS_MODULE,
		.of_match_table = lombo_mtdnor_match,
	},
	.probe  = lombo_nor_probe,
	.remove = lombo_nor_remove,
};

static int __init lombo_mtdnor_driver_init(void)
{
	return spi_register_driver(&lombo_mtdnor_driver);
}
late_initcall(lombo_mtdnor_driver_init);

static void __exit lombo_mtdnor_driver_exit(void)
{
	spi_unregister_driver(&lombo_mtdnor_driver);
}
module_exit(lombo_mtdnor_driver_exit);

MODULE_AUTHOR("lomboswer@lombotech.com>");
MODULE_DESCRIPTION("Lombo MTD SPI driver for nor flash");
MODULE_LICENSE("GPL");
